import sys
import logging
import argparse
from pathlib import Path
from pycocotools.coco import COCO
from pycocoevalcap.eval import COCOEvalCap


class DualOut:
    def __init__(self, *fs):
        self.fs = fs

    def write(self, content):
        for f in self.fs:
            f.write(content)
            f.flush()

    def flush(self):
        for f in self.fs:
            f.flush()


def set_info_logger(args):
    logger = logging.Logger(name="COCO caption evalution", level=logging.INFO)

    eval_info_path = Path(args.caption_log_file)
    formatter = logging.Formatter(fmt="[%(asctime)s]-[%(name)s]-[%(levelname)s]-[%(message)s]")  # formatter for handlers

    file_handler = logging.FileHandler(filename=eval_info_path, mode="a")
    file_handler.setFormatter(formatter)
    file_handler.setLevel(logging.INFO)
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    console_handler.setLevel(logging.INFO)
    logger.addHandler(file_handler); logger.addHandler(console_handler)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--gt-anns-file', type=str, default=None)
    parser.add_argument('--caption-results-file', type=str, default=None)
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    assert Path(args.gt_anns_file).exists(), f"groud truth caption annotations file {args.gt_anns_file} doesn't exist!"
    assert Path(args.caption_results_file).exists(), \
           f"generated caption annotations file {args.caption_results_file} doesn't exist!"
    args.caption_log_file = str(Path(args.caption_results_file).parent / "result_log.txt")
    # set_info_logger(args)  # no use
    log_f = open(Path(args.caption_log_file), "w")
    sys.stdout = DualOut(sys.stdout, log_f)

    coco = COCO(args.gt_anns_file)
    coco_result = coco.loadRes(args.caption_results_file)
    coco_eval = COCOEvalCap(coco, coco_result)
    coco_eval.evaluate()

    sys.stdout = sys.__stdout__
    if not log_f.closed:
        log_f.close()
